共计 3795 个字符,预计需要花费 10 分钟才能阅读完成。
提醒:本文最后更新于 2024-08-30 15:38,文中所关联的信息可能已发生改变,请知悉!
测试代码
import os.path
import logging
import torch
from utils import utils_logger
from utils import utils_image as util
from models.network_rrdbnet import RRDBNet as net
def main():
utils_logger.logger_info('blind_sr_log', log_path='blind_sr_log.log')
logger = logging.getLogger('blind_sr_log')
testsets = 'testsets' # fixed, set path of testsets
testset_Ls = ['RealSRSet'] # ['RealSRSet','DPED']
model_names = ['RRDB','ESRGAN','FSSR_DPED','FSSR_JPEG','RealSR_DPED','RealSR_JPEG']
model_names = ['BSRGAN'] # 'BSRGANx2' for scale factor 2
save_results = True
sf = 4
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
for model_name in model_names:
if model_name in ['BSRGANx2']:
sf = 2
model_path = os.path.join('model_zoo', model_name+'.pth') # set model path
logger.info('{:>16s} : {:s}'.format('Model Name', model_name))
# torch.cuda.set_device(0) # set GPU ID
logger.info('{:>16s} : {:<d}'.format('GPU ID', torch.cuda.current_device()))
torch.cuda.empty_cache()
# --------------------------------
# define network and load model
# --------------------------------
model = net(in_nc=3, out_nc=3, nf=64, nb=23, gc=32, sf=sf) # define network
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
for k, v in model.named_parameters():
v.requires_grad = False
model = model.to(device)
torch.cuda.empty_cache()
for testset_L in testset_Ls:
L_path = os.path.join(testsets, testset_L)
#E_path = os.path.join(testsets, testset_L+'_'+model_name)
E_path = os.path.join(testsets, testset_L+'_results_x'+str(sf))
util.mkdir(E_path)
logger.info('{:>16s} : {:s}'.format('Input Path', L_path))
logger.info('{:>16s} : {:s}'.format('Output Path', E_path))
idx = 0
for img in util.get_image_paths(L_path):
# --------------------------------
# (1) img_L
# --------------------------------
idx += 1
img_name, ext = os.path.splitext(os.path.basename(img))
logger.info('{:->4d} --> {:<s} --> x{:<d}--> {:<s}'.format(idx, model_name, sf, img_name+ext))
img_L = util.imread_uint(img, n_channels=3)
img_L = util.uint2tensor4(img_L)
img_L = img_L.to(device)
# --------------------------------
# (2) inference
# --------------------------------
img_E = model(img_L)
# --------------------------------
# (3) img_E
# --------------------------------
img_E = util.tensor2uint(img_E)
if save_results:
util.imsave(img_E, os.path.join(E_path, img_name+'_'+model_name+'.png'))
if __name__ == '__main__':
main()
结果示例
LogHandlers setup!
22-07-26 21:04:13.379 : Model Name : BSRGAN
22-07-26 21:04:13.386 : GPU ID : 0
[3, 3, 64, 23, 32, 4]
22-07-26 21:04:14.677 : Input Path : testsets\RealSRSet
22-07-26 21:04:14.678 : Output Path : testsets\RealSRSet_results_x4
22-07-26 21:04:14.679 : ---1 --> BSRGAN --> x4--> Lincoln.png
22-07-26 21:04:20.763 : ---2 --> BSRGAN --> x4--> building.png
22-07-26 21:04:21.444 : ---3 --> BSRGAN --> x4--> butterfly.png
22-07-26 21:04:23.191 : ---4 --> BSRGAN --> x4--> butterfly2.png
22-07-26 21:04:23.802 : ---5 --> BSRGAN --> x4--> chip.png
22-07-26 21:04:24.154 : ---6 --> BSRGAN --> x4--> comic1.png
22-07-26 21:04:24.544 : ---7 --> BSRGAN --> x4--> comic2.png
22-07-26 21:04:25.381 : ---8 --> BSRGAN --> x4--> comic3.png
22-07-26 21:04:27.064 : ---9 --> BSRGAN --> x4--> computer.png
22-07-26 21:04:28.882 : --10 --> BSRGAN --> x4--> dog.png
22-07-26 21:04:31.538 : --11 --> BSRGAN --> x4--> dped_crop00061.png
22-07-26 21:04:33.001 : --12 --> BSRGAN --> x4--> foreman.png
22-07-26 21:04:34.198 : --13 --> BSRGAN --> x4--> frog.png
22-07-26 21:04:35.514 : --14 --> BSRGAN --> x4--> oldphoto2.png
22-07-26 21:04:36.157 : --15 --> BSRGAN --> x4--> oldphoto3.png
22-07-26 21:04:36.722 : --16 --> BSRGAN --> x4--> oldphoto6.png
22-07-26 21:04:38.394 : --17 --> BSRGAN --> x4--> painting.png
22-07-26 21:04:39.772 : --18 --> BSRGAN --> x4--> pattern.png
22-07-26 21:04:40.500 : --19 --> BSRGAN --> x4--> ppt3.png
22-07-26 21:04:41.786 : --20 --> BSRGAN --> x4--> tiger.png
部分图片


正文完